{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "CIFAR10_Full_TPU",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "accelerator": "TPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sayakpaul/Training-BatchNorm-and-Only-BatchNorm/blob/master/CIFAR10_Full_TPU.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zRrLNnbuz0fG",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Authenticate yourself to use the TPUs\n",
        "import os\n",
        "\n",
        "IS_COLAB_BACKEND = 'COLAB_GPU' in os.environ  # this is always set on Colab, the value is 0 or 1 depending on GPU presence\n",
        "if IS_COLAB_BACKEND:\n",
        "  from google.colab import auth\n",
        "  # Authenticates the Colab machine and also the TPU using your\n",
        "  # credentials so that they can access your private GCS buckets.\n",
        "  auth.authenticate_user()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "lIYdn1woOS1n",
        "outputId": "a2418388-e953-4d6b-8a0e-a37051745e67",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "# TensorFlow Imports\n",
        "import tensorflow as tf\n",
        "print(tf.__version__)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "2.2.0-rc3\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UpKUNu1Uz3ld",
        "colab_type": "code",
        "outputId": "56ab3e83-5810-4e70-d72e-4cc939fca5d6",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 731
        }
      },
      "source": [
        "# Detect hardware\n",
        "try:\n",
        "  tpu = tf.distribute.cluster_resolver.TPUClusterResolver() # TPU detection\n",
        "except ValueError:\n",
        "  tpu = None\n",
        "  gpus = tf.config.experimental.list_logical_devices(\"GPU\")\n",
        "    \n",
        "# Select appropriate distribution strategy\n",
        "if tpu:\n",
        "  tf.config.experimental_connect_to_cluster(tpu)\n",
        "  tf.tpu.experimental.initialize_tpu_system(tpu)\n",
        "  strategy = tf.distribute.experimental.TPUStrategy(tpu) # Going back and forth between TPU and host is expensive. Better to run 128 batches on the TPU before reporting back.\n",
        "  print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])  \n",
        "elif len(gpus) > 1:\n",
        "  strategy = tf.distribute.MirroredStrategy([gpu.name for gpu in gpus])\n",
        "  print('Running on multiple GPUs ', [gpu.name for gpu in gpus])\n",
        "elif len(gpus) == 1:\n",
        "  strategy = tf.distribute.get_strategy() # default strategy that works on CPU and single GPU\n",
        "  print('Running on single GPU ', gpus[0].name)\n",
        "else:\n",
        "  strategy = tf.distribute.get_strategy() # default strategy that works on CPU and single GPU\n",
        "  print('Running on CPU')\n",
        "print(\"Number of accelerators: \", strategy.num_replicas_in_sync)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Initializing the TPU system: grpc://10.120.151.66:8470\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Initializing the TPU system: grpc://10.120.151.66:8470\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Clearing out eager caches\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Clearing out eager caches\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Finished initializing TPU system.\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Finished initializing TPU system.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Found TPU system:\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Found TPU system:\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Num TPU Cores: 8\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Num TPU Cores: 8\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Num TPU Workers: 1\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Num TPU Workers: 1\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Num TPU Cores Per Worker: 8\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Num TPU Cores Per Worker: 8\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Running on TPU  ['10.120.151.66:8470']\n",
            "Number of accelerators:  8\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "l9tHbooBz95X",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install wandb -q\n",
        "import wandb\n",
        "wandb.login()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5-_4ROoYRHJq",
        "colab_type": "code",
        "outputId": "dba6a6d7-0daa-45dd-d0bf-d0e36841ca6a",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 204
        }
      },
      "source": [
        "!wget https://raw.githubusercontent.com/GoogleCloudPlatform/keras-idiomatic-programmer/master/zoo/resnet/resnet_cifar10.py"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "--2020-05-03 12:26:54--  https://raw.githubusercontent.com/GoogleCloudPlatform/keras-idiomatic-programmer/master/zoo/resnet/resnet_cifar10.py\n",
            "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 151.101.0.133, 151.101.64.133, 151.101.128.133, ...\n",
            "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|151.101.0.133|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 6064 (5.9K) [text/plain]\n",
            "Saving to: ‘resnet_cifar10.py’\n",
            "\n",
            "\rresnet_cifar10.py     0%[                    ]       0  --.-KB/s               \rresnet_cifar10.py   100%[===================>]   5.92K  --.-KB/s    in 0s      \n",
            "\n",
            "2020-05-03 12:26:54 (109 MB/s) - ‘resnet_cifar10.py’ saved [6064/6064]\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "txX7OoEpROgI",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Other imports\n",
        "from tensorflow.keras.layers import *\n",
        "from tensorflow.keras.models import *\n",
        "from wandb.keras import WandbCallback\n",
        "import matplotlib.pyplot as plt\n",
        "import tensorflow as tf\n",
        "import resnet_cifar10\n",
        "import numpy as np\n",
        "import time\n",
        "\n",
        "# Random seed fixation\n",
        "tf.random.set_seed(666)\n",
        "np.random.seed(666)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1NlIHQ-rRWlX",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def get_training_model():\n",
        "    # ResNet20\n",
        "    n = 2\n",
        "    depth =  n * 9 + 2\n",
        "    n_blocks = ((depth - 2) // 9) - 1\n",
        "\n",
        "    # The input tensor\n",
        "    inputs = Input(shape=(32, 32, 3))\n",
        "\n",
        "    # The Stem Convolution Group\n",
        "    x = resnet_cifar10.stem(inputs)\n",
        "\n",
        "    # The learner\n",
        "    x = resnet_cifar10.learner(x, n_blocks)\n",
        "\n",
        "    # The Classifier for 10 classes\n",
        "    outputs = resnet_cifar10.classifier(x, 10)\n",
        "\n",
        "    # Instantiate the Model\n",
        "    model = Model(inputs, outputs)\n",
        "\n",
        "    model.compile(loss=\"sparse_categorical_crossentropy\", optimizer=\"sgd\", metrics=[\"accuracy\"])\n",
        "    \n",
        "    return model"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Jp-VUHKiRqo4",
        "colab_type": "code",
        "outputId": "70025bc3-7abb-4181-ca92-a4077ec00d0b",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        }
      },
      "source": [
        "# Load the training set of CIFAR10\n",
        "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Downloading data from https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz\n",
            "170500096/170498071 [==============================] - 3s 0us/step\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dNhaSmKjR2Vq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "BATCH_SIZE = 128 * strategy.num_replicas_in_sync\n",
        "\n",
        "def normalize(image, label):\n",
        "    image = tf.image.convert_image_dtype(image, tf.float32)\n",
        "    label = tf.cast(label, tf.int32)\n",
        "    return image, label\n",
        "\n",
        "def augment(image,label):\n",
        "    image = tf.image.resize_with_crop_or_pad(image, 40, 40) # Add 8 pixels of padding\n",
        "    image = tf.image.random_crop(image, size=[32, 32, 3]) # Random crop back to 32x32\n",
        "    image = tf.image.random_brightness(image, max_delta=0.5) # Random brightness\n",
        "    image = tf.clip_by_value(image, 0., 1.)\n",
        "\n",
        "    return image, label\n",
        "\n",
        "train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))\n",
        "train_ds = (\n",
        "    train_ds\n",
        "    .shuffle(1024)\n",
        "    .map(normalize, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "    .map(augment, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "    .batch(BATCH_SIZE)\n",
        "    .prefetch(tf.data.experimental.AUTOTUNE)\n",
        ")\n",
        "\n",
        "test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test))\n",
        "test_ds = (\n",
        "    test_ds\n",
        "    .map(normalize, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "    .batch(BATCH_SIZE)\n",
        "    .prefetch(tf.data.experimental.AUTOTUNE)\n",
        ")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ymnw4EirV8HT",
        "colab_type": "code",
        "outputId": "b18d7c76-8fee-472d-efe7-68c59f09b80d",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        }
      },
      "source": [
        "with strategy.scope():\n",
        "    model = get_training_model()\n",
        "model.summary()"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Model: \"model_1\"\n",
            "__________________________________________________________________________________________________\n",
            "Layer (type)                    Output Shape         Param #     Connected to                     \n",
            "==================================================================================================\n",
            "input_2 (InputLayer)            [(None, 32, 32, 3)]  0                                            \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_166 (Conv2D)             (None, 32, 32, 16)   448         input_2[0][0]                    \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_164 (BatchN (None, 32, 32, 16)   64          conv2d_166[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_164 (ReLU)                (None, 32, 32, 16)   0           batch_normalization_164[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_168 (Conv2D)             (None, 32, 32, 16)   272         re_lu_164[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_165 (BatchN (None, 32, 32, 16)   64          conv2d_168[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_165 (ReLU)                (None, 32, 32, 16)   0           batch_normalization_165[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_169 (Conv2D)             (None, 32, 32, 16)   2320        re_lu_165[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_166 (BatchN (None, 32, 32, 16)   64          conv2d_169[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_166 (ReLU)                (None, 32, 32, 16)   0           batch_normalization_166[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_170 (Conv2D)             (None, 32, 32, 64)   1088        re_lu_166[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_167 (Conv2D)             (None, 32, 32, 64)   1088        re_lu_164[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_167 (BatchN (None, 32, 32, 64)   256         conv2d_170[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "add_54 (Add)                    (None, 32, 32, 64)   0           conv2d_167[0][0]                 \n",
            "                                                                 batch_normalization_167[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_167 (ReLU)                (None, 32, 32, 64)   0           add_54[0][0]                     \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_171 (Conv2D)             (None, 32, 32, 16)   1040        re_lu_167[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_168 (BatchN (None, 32, 32, 16)   64          conv2d_171[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_168 (ReLU)                (None, 32, 32, 16)   0           batch_normalization_168[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_172 (Conv2D)             (None, 32, 32, 16)   2320        re_lu_168[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_169 (BatchN (None, 32, 32, 16)   64          conv2d_172[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_169 (ReLU)                (None, 32, 32, 16)   0           batch_normalization_169[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_173 (Conv2D)             (None, 32, 32, 64)   1088        re_lu_169[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_170 (BatchN (None, 32, 32, 64)   256         conv2d_173[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "add_55 (Add)                    (None, 32, 32, 64)   0           batch_normalization_170[0][0]    \n",
            "                                                                 re_lu_167[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_170 (ReLU)                (None, 32, 32, 64)   0           add_55[0][0]                     \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_175 (Conv2D)             (None, 32, 32, 64)   4160        re_lu_170[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_171 (BatchN (None, 32, 32, 64)   256         conv2d_175[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_171 (ReLU)                (None, 32, 32, 64)   0           batch_normalization_171[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_176 (Conv2D)             (None, 16, 16, 64)   36928       re_lu_171[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_172 (BatchN (None, 16, 16, 64)   256         conv2d_176[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_172 (ReLU)                (None, 16, 16, 64)   0           batch_normalization_172[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_177 (Conv2D)             (None, 16, 16, 128)  8320        re_lu_172[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_174 (Conv2D)             (None, 16, 16, 128)  8320        re_lu_170[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_173 (BatchN (None, 16, 16, 128)  512         conv2d_177[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "add_56 (Add)                    (None, 16, 16, 128)  0           conv2d_174[0][0]                 \n",
            "                                                                 batch_normalization_173[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_173 (ReLU)                (None, 16, 16, 128)  0           add_56[0][0]                     \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_178 (Conv2D)             (None, 16, 16, 64)   8256        re_lu_173[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_174 (BatchN (None, 16, 16, 64)   256         conv2d_178[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_174 (ReLU)                (None, 16, 16, 64)   0           batch_normalization_174[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_179 (Conv2D)             (None, 16, 16, 64)   36928       re_lu_174[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_175 (BatchN (None, 16, 16, 64)   256         conv2d_179[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_175 (ReLU)                (None, 16, 16, 64)   0           batch_normalization_175[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_180 (Conv2D)             (None, 16, 16, 128)  8320        re_lu_175[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_176 (BatchN (None, 16, 16, 128)  512         conv2d_180[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "add_57 (Add)                    (None, 16, 16, 128)  0           batch_normalization_176[0][0]    \n",
            "                                                                 re_lu_173[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_176 (ReLU)                (None, 16, 16, 128)  0           add_57[0][0]                     \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_182 (Conv2D)             (None, 16, 16, 128)  16512       re_lu_176[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_177 (BatchN (None, 16, 16, 128)  512         conv2d_182[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_177 (ReLU)                (None, 16, 16, 128)  0           batch_normalization_177[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_183 (Conv2D)             (None, 8, 8, 128)    147584      re_lu_177[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_178 (BatchN (None, 8, 8, 128)    512         conv2d_183[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_178 (ReLU)                (None, 8, 8, 128)    0           batch_normalization_178[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_184 (Conv2D)             (None, 8, 8, 256)    33024       re_lu_178[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_181 (Conv2D)             (None, 8, 8, 256)    33024       re_lu_176[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_179 (BatchN (None, 8, 8, 256)    1024        conv2d_184[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "add_58 (Add)                    (None, 8, 8, 256)    0           conv2d_181[0][0]                 \n",
            "                                                                 batch_normalization_179[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_179 (ReLU)                (None, 8, 8, 256)    0           add_58[0][0]                     \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_185 (Conv2D)             (None, 8, 8, 128)    32896       re_lu_179[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_180 (BatchN (None, 8, 8, 128)    512         conv2d_185[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_180 (ReLU)                (None, 8, 8, 128)    0           batch_normalization_180[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_186 (Conv2D)             (None, 8, 8, 128)    147584      re_lu_180[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_181 (BatchN (None, 8, 8, 128)    512         conv2d_186[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_181 (ReLU)                (None, 8, 8, 128)    0           batch_normalization_181[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_187 (Conv2D)             (None, 8, 8, 256)    33024       re_lu_181[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_182 (BatchN (None, 8, 8, 256)    1024        conv2d_187[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "add_59 (Add)                    (None, 8, 8, 256)    0           batch_normalization_182[0][0]    \n",
            "                                                                 re_lu_179[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_182 (ReLU)                (None, 8, 8, 256)    0           add_59[0][0]                     \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_183 (BatchN (None, 8, 8, 256)    1024        re_lu_182[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_183 (ReLU)                (None, 8, 8, 256)    0           batch_normalization_183[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "average_pooling2d_1 (AveragePoo (None, 1, 1, 256)    0           re_lu_183[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "flatten_1 (Flatten)             (None, 256)          0           average_pooling2d_1[0][0]        \n",
            "__________________________________________________________________________________________________\n",
            "dense_1 (Dense)                 (None, 10)           2570        flatten_1[0][0]                  \n",
            "==================================================================================================\n",
            "Total params: 575,114\n",
            "Trainable params: 571,114\n",
            "Non-trainable params: 4,000\n",
            "__________________________________________________________________________________________________\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lNCYm8J2WiMh",
        "colab_type": "text"
      },
      "source": [
        "- Total params: 575,114\n",
        "- **Trainable params: 571,114**\n",
        "- Non-trainable params: 4,000\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WOv6rFsbR9z2",
        "colab_type": "code",
        "outputId": "6e3fb4ba-e653-453f-e554-087d280a647e",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        }
      },
      "source": [
        "# Train model\n",
        "wandb.init(project=\"training-bn-only\", id=\"vanilla_resnet\")\n",
        "\n",
        "with strategy.scope():\n",
        "    model = get_training_model()\n",
        "    \n",
        "start = time.time()\n",
        "h = model.fit(train_ds,\n",
        "         validation_data=test_ds,\n",
        "         epochs=75,\n",
        "         callbacks=[WandbCallback()])\n",
        "end = time.time()\n",
        "wandb.log({\"training_time\": end - start})\n",
        "print(\"Network takes {:.3f} seconds to train\".format(end - start))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "                Logging results to <a href=\"https://wandb.com\" target=\"_blank\">Weights & Biases</a> <a href=\"https://docs.wandb.com/integrations/jupyter.html\" target=\"_blank\">(Documentation)</a>.<br/>\n",
              "                Project page: <a href=\"https://app.wandb.ai/sayakpaul/training-bn-only\" target=\"_blank\">https://app.wandb.ai/sayakpaul/training-bn-only</a><br/>\n",
              "                Run page: <a href=\"https://app.wandb.ai/sayakpaul/training-bn-only/runs/vanilla_resnet\" target=\"_blank\">https://app.wandb.ai/sayakpaul/training-bn-only/runs/vanilla_resnet</a><br/>\n",
              "            "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "Epoch 1/75\n",
            "49/49 [==============================] - 10s 208ms/step - loss: 2.3432 - accuracy: 0.1359 - val_loss: 2.3748 - val_accuracy: 0.1517\n",
            "Epoch 2/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 2.2274 - accuracy: 0.1759 - val_loss: 2.2666 - val_accuracy: 0.1680\n",
            "Epoch 3/75\n",
            "49/49 [==============================] - 5s 92ms/step - loss: 2.1778 - accuracy: 0.1971 - val_loss: 2.2408 - val_accuracy: 0.1646\n",
            "Epoch 4/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 2.1347 - accuracy: 0.2145 - val_loss: 2.2012 - val_accuracy: 0.1790\n",
            "Epoch 5/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 2.0963 - accuracy: 0.2235 - val_loss: 2.1505 - val_accuracy: 0.1908\n",
            "Epoch 6/75\n",
            "49/49 [==============================] - 5s 99ms/step - loss: 2.0715 - accuracy: 0.2345 - val_loss: 2.0890 - val_accuracy: 0.2148\n",
            "Epoch 7/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 2.0469 - accuracy: 0.2429 - val_loss: 2.0506 - val_accuracy: 0.2346\n",
            "Epoch 8/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 2.0313 - accuracy: 0.2474 - val_loss: 2.0263 - val_accuracy: 0.2412\n",
            "Epoch 9/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 2.0140 - accuracy: 0.2547 - val_loss: 2.0080 - val_accuracy: 0.2500\n",
            "Epoch 10/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 2.0004 - accuracy: 0.2598 - val_loss: 1.9843 - val_accuracy: 0.2584\n",
            "Epoch 11/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.9883 - accuracy: 0.2652 - val_loss: 1.9715 - val_accuracy: 0.2626\n",
            "Epoch 12/75\n",
            "49/49 [==============================] - 5s 103ms/step - loss: 1.9767 - accuracy: 0.2683 - val_loss: 1.9611 - val_accuracy: 0.2686\n",
            "Epoch 13/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 1.9672 - accuracy: 0.2711 - val_loss: 1.9484 - val_accuracy: 0.2732\n",
            "Epoch 14/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 1.9568 - accuracy: 0.2782 - val_loss: 1.9356 - val_accuracy: 0.2808\n",
            "Epoch 15/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.9469 - accuracy: 0.2813 - val_loss: 1.9279 - val_accuracy: 0.2853\n",
            "Epoch 16/75\n",
            "49/49 [==============================] - 5s 96ms/step - loss: 1.9357 - accuracy: 0.2868 - val_loss: 1.9187 - val_accuracy: 0.2841\n",
            "Epoch 17/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.9261 - accuracy: 0.2892 - val_loss: 1.9065 - val_accuracy: 0.2895\n",
            "Epoch 18/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.9168 - accuracy: 0.2931 - val_loss: 1.9014 - val_accuracy: 0.2947\n",
            "Epoch 19/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.9093 - accuracy: 0.2947 - val_loss: 1.8902 - val_accuracy: 0.2968\n",
            "Epoch 20/75\n",
            "49/49 [==============================] - 5s 104ms/step - loss: 1.9022 - accuracy: 0.3012 - val_loss: 1.8760 - val_accuracy: 0.3031\n",
            "Epoch 21/75\n",
            "49/49 [==============================] - 5s 96ms/step - loss: 1.8937 - accuracy: 0.3020 - val_loss: 1.8699 - val_accuracy: 0.3042\n",
            "Epoch 22/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.8823 - accuracy: 0.3096 - val_loss: 1.8585 - val_accuracy: 0.3077\n",
            "Epoch 23/75\n",
            "49/49 [==============================] - 5s 96ms/step - loss: 1.8739 - accuracy: 0.3092 - val_loss: 1.8480 - val_accuracy: 0.3112\n",
            "Epoch 24/75\n",
            "49/49 [==============================] - 5s 97ms/step - loss: 1.8677 - accuracy: 0.3128 - val_loss: 1.8395 - val_accuracy: 0.3188\n",
            "Epoch 25/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.8596 - accuracy: 0.3159 - val_loss: 1.8274 - val_accuracy: 0.3231\n",
            "Epoch 26/75\n",
            "49/49 [==============================] - 5s 96ms/step - loss: 1.8526 - accuracy: 0.3195 - val_loss: 1.8204 - val_accuracy: 0.3247\n",
            "Epoch 27/75\n",
            "49/49 [==============================] - 5s 97ms/step - loss: 1.8422 - accuracy: 0.3208 - val_loss: 1.8056 - val_accuracy: 0.3313\n",
            "Epoch 28/75\n",
            "49/49 [==============================] - 5s 107ms/step - loss: 1.8353 - accuracy: 0.3239 - val_loss: 1.8017 - val_accuracy: 0.3297\n",
            "Epoch 29/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.8287 - accuracy: 0.3290 - val_loss: 1.7929 - val_accuracy: 0.3290\n",
            "Epoch 30/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 1.8219 - accuracy: 0.3298 - val_loss: 1.7892 - val_accuracy: 0.3299\n",
            "Epoch 31/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.8127 - accuracy: 0.3318 - val_loss: 1.7714 - val_accuracy: 0.3371\n",
            "Epoch 32/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.8042 - accuracy: 0.3337 - val_loss: 1.7641 - val_accuracy: 0.3384\n",
            "Epoch 33/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.7990 - accuracy: 0.3374 - val_loss: 1.7513 - val_accuracy: 0.3441\n",
            "Epoch 34/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.7878 - accuracy: 0.3394 - val_loss: 1.7416 - val_accuracy: 0.3506\n",
            "Epoch 35/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.7816 - accuracy: 0.3435 - val_loss: 1.7364 - val_accuracy: 0.3502\n",
            "Epoch 36/75\n",
            "49/49 [==============================] - 5s 96ms/step - loss: 1.7747 - accuracy: 0.3469 - val_loss: 1.7280 - val_accuracy: 0.3551\n",
            "Epoch 37/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 1.7657 - accuracy: 0.3474 - val_loss: 1.7190 - val_accuracy: 0.3597\n",
            "Epoch 38/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 1.7607 - accuracy: 0.3497 - val_loss: 1.7074 - val_accuracy: 0.3611\n",
            "Epoch 39/75\n",
            "49/49 [==============================] - 5s 109ms/step - loss: 1.7528 - accuracy: 0.3549 - val_loss: 1.7004 - val_accuracy: 0.3641\n",
            "Epoch 40/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.7496 - accuracy: 0.3531 - val_loss: 1.6945 - val_accuracy: 0.3644\n",
            "Epoch 41/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.7397 - accuracy: 0.3597 - val_loss: 1.6878 - val_accuracy: 0.3670\n",
            "Epoch 42/75\n",
            "49/49 [==============================] - 4s 81ms/step - loss: 1.7380 - accuracy: 0.3564 - val_loss: 1.6904 - val_accuracy: 0.3661\n",
            "Epoch 43/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.7294 - accuracy: 0.3614 - val_loss: 1.6695 - val_accuracy: 0.3751\n",
            "Epoch 44/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.7246 - accuracy: 0.3648 - val_loss: 1.6594 - val_accuracy: 0.3811\n",
            "Epoch 45/75\n",
            "49/49 [==============================] - 5s 92ms/step - loss: 1.7196 - accuracy: 0.3663 - val_loss: 1.6586 - val_accuracy: 0.3804\n",
            "Epoch 46/75\n",
            "49/49 [==============================] - 5s 92ms/step - loss: 1.7155 - accuracy: 0.3685 - val_loss: 1.6533 - val_accuracy: 0.3815\n",
            "Epoch 47/75\n",
            "49/49 [==============================] - 5s 96ms/step - loss: 1.7097 - accuracy: 0.3686 - val_loss: 1.6435 - val_accuracy: 0.3889\n",
            "Epoch 48/75\n",
            "49/49 [==============================] - 4s 81ms/step - loss: 1.7051 - accuracy: 0.3706 - val_loss: 1.6488 - val_accuracy: 0.3810\n",
            "Epoch 49/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.6995 - accuracy: 0.3742 - val_loss: 1.6361 - val_accuracy: 0.3880\n",
            "Epoch 50/75\n",
            "49/49 [==============================] - 5s 92ms/step - loss: 1.6957 - accuracy: 0.3745 - val_loss: 1.6286 - val_accuracy: 0.3913\n",
            "Epoch 51/75\n",
            "49/49 [==============================] - 5s 110ms/step - loss: 1.6867 - accuracy: 0.3783 - val_loss: 1.6239 - val_accuracy: 0.3937\n",
            "Epoch 52/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.6835 - accuracy: 0.3787 - val_loss: 1.6201 - val_accuracy: 0.3962\n",
            "Epoch 53/75\n",
            "49/49 [==============================] - 4s 82ms/step - loss: 1.6820 - accuracy: 0.3834 - val_loss: 1.6222 - val_accuracy: 0.3932\n",
            "Epoch 54/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.6766 - accuracy: 0.3827 - val_loss: 1.6166 - val_accuracy: 0.3988\n",
            "Epoch 55/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.6759 - accuracy: 0.3794 - val_loss: 1.6099 - val_accuracy: 0.3978\n",
            "Epoch 56/75\n",
            "49/49 [==============================] - 5s 97ms/step - loss: 1.6667 - accuracy: 0.3849 - val_loss: 1.6032 - val_accuracy: 0.4047\n",
            "Epoch 57/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.6653 - accuracy: 0.3843 - val_loss: 1.5941 - val_accuracy: 0.4072\n",
            "Epoch 58/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 1.6627 - accuracy: 0.3867 - val_loss: 1.5880 - val_accuracy: 0.4060\n",
            "Epoch 59/75\n",
            "49/49 [==============================] - 4s 82ms/step - loss: 1.6534 - accuracy: 0.3891 - val_loss: 1.6193 - val_accuracy: 0.3997\n",
            "Epoch 60/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.6519 - accuracy: 0.3915 - val_loss: 1.5797 - val_accuracy: 0.4079\n",
            "Epoch 61/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.6502 - accuracy: 0.3947 - val_loss: 1.5786 - val_accuracy: 0.4107\n",
            "Epoch 62/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 1.6460 - accuracy: 0.3956 - val_loss: 1.5712 - val_accuracy: 0.4162\n",
            "Epoch 63/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.6392 - accuracy: 0.3948 - val_loss: 1.5639 - val_accuracy: 0.4203\n",
            "Epoch 64/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.6361 - accuracy: 0.4012 - val_loss: 1.5604 - val_accuracy: 0.4219\n",
            "Epoch 65/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 1.6317 - accuracy: 0.3999 - val_loss: 1.5562 - val_accuracy: 0.4236\n",
            "Epoch 66/75\n",
            "49/49 [==============================] - 6s 114ms/step - loss: 1.6308 - accuracy: 0.4021 - val_loss: 1.5557 - val_accuracy: 0.4209\n",
            "Epoch 67/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.6230 - accuracy: 0.4032 - val_loss: 1.5512 - val_accuracy: 0.4236\n",
            "Epoch 68/75\n",
            "49/49 [==============================] - 5s 92ms/step - loss: 1.6233 - accuracy: 0.4025 - val_loss: 1.5452 - val_accuracy: 0.4250\n",
            "Epoch 69/75\n",
            "49/49 [==============================] - 4s 92ms/step - loss: 1.6182 - accuracy: 0.4092 - val_loss: 1.5424 - val_accuracy: 0.4277\n",
            "Epoch 70/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.6159 - accuracy: 0.4069 - val_loss: 1.5349 - val_accuracy: 0.4309\n",
            "Epoch 71/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 1.6115 - accuracy: 0.4103 - val_loss: 1.5341 - val_accuracy: 0.4307\n",
            "Epoch 72/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.6073 - accuracy: 0.4099 - val_loss: 1.5331 - val_accuracy: 0.4324\n",
            "Epoch 73/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 1.6026 - accuracy: 0.4139 - val_loss: 1.5272 - val_accuracy: 0.4352\n",
            "Epoch 74/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.6011 - accuracy: 0.4124 - val_loss: 1.5262 - val_accuracy: 0.4372\n",
            "Epoch 75/75\n",
            "49/49 [==============================] - 5s 96ms/step - loss: 1.5972 - accuracy: 0.4144 - val_loss: 1.5180 - val_accuracy: 0.4380\n",
            "Network takes 371.131 seconds to train\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NXBgHyMLggnE",
        "colab_type": "code",
        "outputId": "92cc3c5f-bc10-4377-9933-6ac8e00162f1",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 119
        }
      },
      "source": [
        "# Save the weights to GCS bucket for reproducibility\n",
        "PROJECT_ID = \"fast-ai-exploration\" \n",
        "!gcloud config set project $PROJECT_ID"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Updated property [core/project].\n",
            "\n",
            "\n",
            "To take a quick anonymous survey, run:\n",
            "  $ gcloud survey\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AT1eEEDXiFKT",
        "colab_type": "code",
        "outputId": "6f12aafc-2067-4b28-bad4-6db2556054b3",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "!gsutil mb gs://batch_norm_tpu"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Creating gs://batch_norm_tpu/...\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Y5g9rjAjhWua",
        "colab_type": "code",
        "outputId": "2719b404-bbf5-4a09-9179-04663964adbf",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 173
        }
      },
      "source": [
        "model.save(\"gs://batch_norm_tpu/vanilla_resnet\")\n",
        "# !gsutil -m cp -r vanilla_resnet gs://batch_norm_tpu"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/resource_variable_ops.py:1817: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "If using Keras pass *_constraint arguments to layers.\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/resource_variable_ops.py:1817: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "If using Keras pass *_constraint arguments to layers.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Assets written to: gs://batch_norm_tpu/vanilla_resnet/assets\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Assets written to: gs://batch_norm_tpu/vanilla_resnet/assets\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kTwu12AuBHzM",
        "colab_type": "code",
        "outputId": "e7cb4b2b-8f3b-4331-e369-155b5366df14",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 102
        }
      },
      "source": [
        "!gsutil ls -lh gs://batch_norm_tpu/vanilla_resnet/"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "       0 B  2020-05-03T12:35:08Z  gs://batch_norm_tpu/vanilla_resnet/\n",
            "  4.94 MiB  2020-05-03T12:35:11Z  gs://batch_norm_tpu/vanilla_resnet/saved_model.pb\n",
            "                                 gs://batch_norm_tpu/vanilla_resnet/assets/\n",
            "                                 gs://batch_norm_tpu/vanilla_resnet/variables/\n",
            "TOTAL: 2 objects, 5179967 bytes (4.94 MiB)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rM5mFUHkSWcp",
        "colab_type": "code",
        "outputId": "accc78f6-958c-409f-ce50-eb6543d6db6d",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        }
      },
      "source": [
        "# Train model with a decay schedule\n",
        "wandb.init(project=\"training-bn-only\", id=\"resnet_lrsch_1\")\n",
        "\n",
        "first_decay_steps = 1000\n",
        "lr_decayed_fn = (\n",
        "  tf.keras.experimental.CosineDecay(\n",
        "      initial_learning_rate=0.1,\n",
        "      decay_steps=first_decay_steps))\n",
        "\n",
        "with strategy.scope():\n",
        "    model = get_training_model()\n",
        "start = time.time()\n",
        "h = model.fit(train_ds,\n",
        "         validation_data=test_ds,\n",
        "         epochs=75,\n",
        "         callbacks=[WandbCallback()])\n",
        "end = time.time()\n",
        "wandb.log({\"training_time\": end - start})\n",
        "print(\"Network takes {:.3f} seconds to train\".format(end - start))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "                Logging results to <a href=\"https://wandb.com\" target=\"_blank\">Weights & Biases</a> <a href=\"https://docs.wandb.com/integrations/jupyter.html\" target=\"_blank\">(Documentation)</a>.<br/>\n",
              "                Project page: <a href=\"https://app.wandb.ai/sayakpaul/training-bn-only\" target=\"_blank\">https://app.wandb.ai/sayakpaul/training-bn-only</a><br/>\n",
              "                Run page: <a href=\"https://app.wandb.ai/sayakpaul/training-bn-only/runs/resnet_lrsch_1\" target=\"_blank\">https://app.wandb.ai/sayakpaul/training-bn-only/runs/resnet_lrsch_1</a><br/>\n",
              "            "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "Epoch 1/75\n",
            "49/49 [==============================] - 10s 213ms/step - loss: 2.3547 - accuracy: 0.1189 - val_loss: 2.3836 - val_accuracy: 0.1046\n",
            "Epoch 2/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 2.2579 - accuracy: 0.1552 - val_loss: 2.2903 - val_accuracy: 0.1353\n",
            "Epoch 3/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 2.2239 - accuracy: 0.1825 - val_loss: 2.2768 - val_accuracy: 0.1392\n",
            "Epoch 4/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 2.1904 - accuracy: 0.1970 - val_loss: 2.2526 - val_accuracy: 0.1461\n",
            "Epoch 5/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 2.1583 - accuracy: 0.2121 - val_loss: 2.2208 - val_accuracy: 0.1580\n",
            "Epoch 6/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 2.1250 - accuracy: 0.2253 - val_loss: 2.1772 - val_accuracy: 0.1773\n",
            "Epoch 7/75\n",
            "49/49 [==============================] - 5s 96ms/step - loss: 2.0902 - accuracy: 0.2371 - val_loss: 2.1259 - val_accuracy: 0.2020\n",
            "Epoch 8/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 2.0589 - accuracy: 0.2505 - val_loss: 2.0778 - val_accuracy: 0.2150\n",
            "Epoch 9/75\n",
            "49/49 [==============================] - 5s 96ms/step - loss: 2.0280 - accuracy: 0.2577 - val_loss: 2.0310 - val_accuracy: 0.2353\n",
            "Epoch 10/75\n",
            "49/49 [==============================] - 5s 96ms/step - loss: 1.9999 - accuracy: 0.2652 - val_loss: 1.9872 - val_accuracy: 0.2540\n",
            "Epoch 11/75\n",
            "49/49 [==============================] - 7s 134ms/step - loss: 1.9731 - accuracy: 0.2730 - val_loss: 1.9487 - val_accuracy: 0.2708\n",
            "Epoch 12/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.9493 - accuracy: 0.2827 - val_loss: 1.9150 - val_accuracy: 0.2820\n",
            "Epoch 13/75\n",
            "49/49 [==============================] - 5s 96ms/step - loss: 1.9294 - accuracy: 0.2873 - val_loss: 1.8932 - val_accuracy: 0.2903\n",
            "Epoch 14/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.9121 - accuracy: 0.2897 - val_loss: 1.8729 - val_accuracy: 0.2975\n",
            "Epoch 15/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.8910 - accuracy: 0.2991 - val_loss: 1.8591 - val_accuracy: 0.3054\n",
            "Epoch 16/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.8796 - accuracy: 0.3014 - val_loss: 1.8381 - val_accuracy: 0.3117\n",
            "Epoch 17/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.8668 - accuracy: 0.3070 - val_loss: 1.8172 - val_accuracy: 0.3192\n",
            "Epoch 18/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.8551 - accuracy: 0.3094 - val_loss: 1.8040 - val_accuracy: 0.3229\n",
            "Epoch 19/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.8416 - accuracy: 0.3161 - val_loss: 1.7968 - val_accuracy: 0.3277\n",
            "Epoch 20/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.8313 - accuracy: 0.3171 - val_loss: 1.7813 - val_accuracy: 0.3372\n",
            "Epoch 21/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.8218 - accuracy: 0.3216 - val_loss: 1.7663 - val_accuracy: 0.3442\n",
            "Epoch 22/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.8126 - accuracy: 0.3259 - val_loss: 1.7523 - val_accuracy: 0.3542\n",
            "Epoch 23/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.8042 - accuracy: 0.3273 - val_loss: 1.7398 - val_accuracy: 0.3544\n",
            "Epoch 24/75\n",
            "49/49 [==============================] - 4s 83ms/step - loss: 1.7931 - accuracy: 0.3317 - val_loss: 1.7619 - val_accuracy: 0.3406\n",
            "Epoch 25/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.7854 - accuracy: 0.3371 - val_loss: 1.7365 - val_accuracy: 0.3531\n",
            "Epoch 26/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.7783 - accuracy: 0.3364 - val_loss: 1.7179 - val_accuracy: 0.3612\n",
            "Epoch 27/75\n",
            "49/49 [==============================] - 5s 96ms/step - loss: 1.7727 - accuracy: 0.3393 - val_loss: 1.7164 - val_accuracy: 0.3630\n",
            "Epoch 28/75\n",
            "49/49 [==============================] - 4s 83ms/step - loss: 1.7582 - accuracy: 0.3453 - val_loss: 1.7539 - val_accuracy: 0.3431\n",
            "Epoch 29/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.7605 - accuracy: 0.3420 - val_loss: 1.6972 - val_accuracy: 0.3723\n",
            "Epoch 30/75\n",
            "49/49 [==============================] - 5s 96ms/step - loss: 1.7512 - accuracy: 0.3473 - val_loss: 1.6938 - val_accuracy: 0.3687\n",
            "Epoch 31/75\n",
            "49/49 [==============================] - 5s 97ms/step - loss: 1.7439 - accuracy: 0.3511 - val_loss: 1.6737 - val_accuracy: 0.3823\n",
            "Epoch 32/75\n",
            "49/49 [==============================] - 4s 82ms/step - loss: 1.7410 - accuracy: 0.3529 - val_loss: 1.6775 - val_accuracy: 0.3755\n",
            "Epoch 33/75\n",
            "49/49 [==============================] - 6s 125ms/step - loss: 1.7354 - accuracy: 0.3538 - val_loss: 1.6722 - val_accuracy: 0.3824\n",
            "Epoch 34/75\n",
            "49/49 [==============================] - 5s 96ms/step - loss: 1.7267 - accuracy: 0.3598 - val_loss: 1.6566 - val_accuracy: 0.3876\n",
            "Epoch 35/75\n",
            "49/49 [==============================] - 4s 82ms/step - loss: 1.7232 - accuracy: 0.3601 - val_loss: 1.6617 - val_accuracy: 0.3861\n",
            "Epoch 36/75\n",
            "49/49 [==============================] - 4s 80ms/step - loss: 1.7161 - accuracy: 0.3616 - val_loss: 1.6749 - val_accuracy: 0.3855\n",
            "Epoch 37/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.7121 - accuracy: 0.3624 - val_loss: 1.6382 - val_accuracy: 0.3964\n",
            "Epoch 38/75\n",
            "49/49 [==============================] - 4s 82ms/step - loss: 1.7090 - accuracy: 0.3644 - val_loss: 1.6720 - val_accuracy: 0.3816\n",
            "Epoch 39/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.7030 - accuracy: 0.3672 - val_loss: 1.6323 - val_accuracy: 0.3973\n",
            "Epoch 40/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 1.6965 - accuracy: 0.3711 - val_loss: 1.6250 - val_accuracy: 0.4038\n",
            "Epoch 41/75\n",
            "49/49 [==============================] - 4s 87ms/step - loss: 1.6947 - accuracy: 0.3724 - val_loss: 1.6261 - val_accuracy: 0.4008\n",
            "Epoch 42/75\n",
            "49/49 [==============================] - 5s 97ms/step - loss: 1.6929 - accuracy: 0.3704 - val_loss: 1.6186 - val_accuracy: 0.4040\n",
            "Epoch 43/75\n",
            "49/49 [==============================] - 4s 85ms/step - loss: 1.6834 - accuracy: 0.3761 - val_loss: 1.7075 - val_accuracy: 0.3808\n",
            "Epoch 44/75\n",
            "49/49 [==============================] - 4s 83ms/step - loss: 1.6830 - accuracy: 0.3773 - val_loss: 1.6563 - val_accuracy: 0.3924\n",
            "Epoch 45/75\n",
            "49/49 [==============================] - 4s 83ms/step - loss: 1.6775 - accuracy: 0.3779 - val_loss: 1.6296 - val_accuracy: 0.3939\n",
            "Epoch 46/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.6759 - accuracy: 0.3772 - val_loss: 1.6021 - val_accuracy: 0.4141\n",
            "Epoch 47/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.6697 - accuracy: 0.3812 - val_loss: 1.5990 - val_accuracy: 0.4055\n",
            "Epoch 48/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.6642 - accuracy: 0.3835 - val_loss: 1.5938 - val_accuracy: 0.4131\n",
            "Epoch 49/75\n",
            "49/49 [==============================] - 4s 84ms/step - loss: 1.6618 - accuracy: 0.3838 - val_loss: 1.5991 - val_accuracy: 0.4135\n",
            "Epoch 50/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.6563 - accuracy: 0.3864 - val_loss: 1.5849 - val_accuracy: 0.4191\n",
            "Epoch 51/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.6539 - accuracy: 0.3883 - val_loss: 1.5794 - val_accuracy: 0.4159\n",
            "Epoch 52/75\n",
            "49/49 [==============================] - 4s 82ms/step - loss: 1.6512 - accuracy: 0.3901 - val_loss: 1.5825 - val_accuracy: 0.4150\n",
            "Epoch 53/75\n",
            "49/49 [==============================] - 5s 96ms/step - loss: 1.6491 - accuracy: 0.3918 - val_loss: 1.5675 - val_accuracy: 0.4246\n",
            "Epoch 54/75\n",
            "49/49 [==============================] - 4s 85ms/step - loss: 1.6434 - accuracy: 0.3915 - val_loss: 1.5918 - val_accuracy: 0.4158\n",
            "Epoch 55/75\n",
            "49/49 [==============================] - 4s 83ms/step - loss: 1.6432 - accuracy: 0.3933 - val_loss: 1.5717 - val_accuracy: 0.4198\n",
            "Epoch 56/75\n",
            "49/49 [==============================] - 4s 84ms/step - loss: 1.6369 - accuracy: 0.3962 - val_loss: 1.5769 - val_accuracy: 0.4214\n",
            "Epoch 57/75\n",
            "49/49 [==============================] - 5s 96ms/step - loss: 1.6335 - accuracy: 0.3955 - val_loss: 1.5649 - val_accuracy: 0.4247\n",
            "Epoch 58/75\n",
            "49/49 [==============================] - 4s 83ms/step - loss: 1.6279 - accuracy: 0.3968 - val_loss: 1.5875 - val_accuracy: 0.4182\n",
            "Epoch 59/75\n",
            "49/49 [==============================] - 6s 132ms/step - loss: 1.6248 - accuracy: 0.4006 - val_loss: 1.5500 - val_accuracy: 0.4289\n",
            "Epoch 60/75\n",
            "49/49 [==============================] - 5s 98ms/step - loss: 1.6241 - accuracy: 0.3982 - val_loss: 1.5424 - val_accuracy: 0.4342\n",
            "Epoch 61/75\n",
            "49/49 [==============================] - 4s 84ms/step - loss: 1.6207 - accuracy: 0.4026 - val_loss: 1.5439 - val_accuracy: 0.4293\n",
            "Epoch 62/75\n",
            "49/49 [==============================] - 4s 83ms/step - loss: 1.6166 - accuracy: 0.4058 - val_loss: 1.5812 - val_accuracy: 0.4235\n",
            "Epoch 63/75\n",
            "49/49 [==============================] - 4s 82ms/step - loss: 1.6135 - accuracy: 0.4051 - val_loss: 1.5642 - val_accuracy: 0.4243\n",
            "Epoch 64/75\n",
            "49/49 [==============================] - 4s 82ms/step - loss: 1.6107 - accuracy: 0.4034 - val_loss: 1.5437 - val_accuracy: 0.4308\n",
            "Epoch 65/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.6086 - accuracy: 0.4069 - val_loss: 1.5278 - val_accuracy: 0.4365\n",
            "Epoch 66/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.6032 - accuracy: 0.4072 - val_loss: 1.5203 - val_accuracy: 0.4407\n",
            "Epoch 67/75\n",
            "49/49 [==============================] - 4s 80ms/step - loss: 1.6038 - accuracy: 0.4092 - val_loss: 1.6356 - val_accuracy: 0.4119\n",
            "Epoch 68/75\n",
            "49/49 [==============================] - 4s 81ms/step - loss: 1.6001 - accuracy: 0.4090 - val_loss: 1.5343 - val_accuracy: 0.4336\n",
            "Epoch 69/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.5940 - accuracy: 0.4136 - val_loss: 1.5150 - val_accuracy: 0.4483\n",
            "Epoch 70/75\n",
            "49/49 [==============================] - 4s 81ms/step - loss: 1.5913 - accuracy: 0.4146 - val_loss: 1.5209 - val_accuracy: 0.4400\n",
            "Epoch 71/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.5903 - accuracy: 0.4138 - val_loss: 1.5149 - val_accuracy: 0.4474\n",
            "Epoch 72/75\n",
            "49/49 [==============================] - 4s 83ms/step - loss: 1.5876 - accuracy: 0.4141 - val_loss: 1.6278 - val_accuracy: 0.4163\n",
            "Epoch 73/75\n",
            "49/49 [==============================] - 4s 82ms/step - loss: 1.5829 - accuracy: 0.4192 - val_loss: 1.5575 - val_accuracy: 0.4318\n",
            "Epoch 74/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 1.5826 - accuracy: 0.4198 - val_loss: 1.5026 - val_accuracy: 0.4503\n",
            "Epoch 75/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.5747 - accuracy: 0.4221 - val_loss: 1.4982 - val_accuracy: 0.4515\n",
            "Network takes 362.584 seconds to train\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5o1kTTj8iUL3",
        "colab_type": "code",
        "outputId": "da31d73a-d4bd-4da0-de0c-1d4d854fb182",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 68
        }
      },
      "source": [
        "model.save_weights(\"resnet_lrsch_1.h5\")\n",
        "!gsutil -m cp resnet_lrsch_1.h5 gs://batch_norm_tpu/"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Copying file://resnet_lrsch_1.h5 [Content-Type=application/octet-stream]...\n",
            "/ [1/1 files][  2.4 MiB/  2.4 MiB] 100% Done                                    \n",
            "Operation completed over 1 objects/2.4 MiB.                                      \n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SLFNETnFZL_J",
        "colab_type": "code",
        "outputId": "f1deb0fb-fa5c-466c-9a87-7c98238227e0",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        }
      },
      "source": [
        "# Train model with a decay schedule\n",
        "wandb.init(project=\"training-bn-only\", id=\"resnet_lrsch_2\")\n",
        "\n",
        "first_decay_steps = 1000\n",
        "lr_decayed_fn = (\n",
        "  tf.keras.experimental.LinearCosineDecay(\n",
        "      initial_learning_rate=0.1,\n",
        "      decay_steps=first_decay_steps))\n",
        "\n",
        "with strategy.scope():\n",
        "    model = get_training_model()\n",
        "start = time.time()\n",
        "h = model.fit(train_ds,\n",
        "         validation_data=test_ds,\n",
        "         epochs=75,\n",
        "         callbacks=[WandbCallback()])\n",
        "end = time.time()\n",
        "wandb.log({\"training_time\": end - start})\n",
        "print(\"Network takes {:.3f} seconds to train\".format(end - start))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "                Logging results to <a href=\"https://wandb.com\" target=\"_blank\">Weights & Biases</a> <a href=\"https://docs.wandb.com/integrations/jupyter.html\" target=\"_blank\">(Documentation)</a>.<br/>\n",
              "                Project page: <a href=\"https://app.wandb.ai/sayakpaul/training-bn-only\" target=\"_blank\">https://app.wandb.ai/sayakpaul/training-bn-only</a><br/>\n",
              "                Run page: <a href=\"https://app.wandb.ai/sayakpaul/training-bn-only/runs/resnet_lrsch_2\" target=\"_blank\">https://app.wandb.ai/sayakpaul/training-bn-only/runs/resnet_lrsch_2</a><br/>\n",
              "            "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "Epoch 1/75\n",
            "49/49 [==============================] - 11s 217ms/step - loss: 2.3147 - accuracy: 0.1224 - val_loss: 2.2760 - val_accuracy: 0.1195\n",
            "Epoch 2/75\n",
            "49/49 [==============================] - 5s 96ms/step - loss: 2.2221 - accuracy: 0.1673 - val_loss: 2.2651 - val_accuracy: 0.1324\n",
            "Epoch 3/75\n",
            "49/49 [==============================] - 5s 96ms/step - loss: 2.1814 - accuracy: 0.1892 - val_loss: 2.2502 - val_accuracy: 0.1333\n",
            "Epoch 4/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 2.1487 - accuracy: 0.2052 - val_loss: 2.2196 - val_accuracy: 0.1400\n",
            "Epoch 5/75\n",
            "49/49 [==============================] - 5s 96ms/step - loss: 2.1181 - accuracy: 0.2160 - val_loss: 2.1744 - val_accuracy: 0.1585\n",
            "Epoch 6/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 2.0934 - accuracy: 0.2263 - val_loss: 2.1275 - val_accuracy: 0.1840\n",
            "Epoch 7/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 2.0672 - accuracy: 0.2372 - val_loss: 2.0858 - val_accuracy: 0.2020\n",
            "Epoch 8/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 2.0447 - accuracy: 0.2455 - val_loss: 2.0424 - val_accuracy: 0.2249\n",
            "Epoch 9/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 2.0255 - accuracy: 0.2550 - val_loss: 1.9981 - val_accuracy: 0.2470\n",
            "Epoch 10/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 2.0056 - accuracy: 0.2622 - val_loss: 1.9645 - val_accuracy: 0.2606\n",
            "Epoch 11/75\n",
            "49/49 [==============================] - 7s 142ms/step - loss: 1.9865 - accuracy: 0.2695 - val_loss: 1.9458 - val_accuracy: 0.2678\n",
            "Epoch 12/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.9727 - accuracy: 0.2754 - val_loss: 1.9244 - val_accuracy: 0.2775\n",
            "Epoch 13/75\n",
            "49/49 [==============================] - 5s 98ms/step - loss: 1.9563 - accuracy: 0.2806 - val_loss: 1.9065 - val_accuracy: 0.2863\n",
            "Epoch 14/75\n",
            "49/49 [==============================] - 5s 96ms/step - loss: 1.9417 - accuracy: 0.2838 - val_loss: 1.8923 - val_accuracy: 0.2933\n",
            "Epoch 15/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.9305 - accuracy: 0.2867 - val_loss: 1.8797 - val_accuracy: 0.2978\n",
            "Epoch 16/75\n",
            "49/49 [==============================] - 5s 96ms/step - loss: 1.9164 - accuracy: 0.2937 - val_loss: 1.8699 - val_accuracy: 0.3008\n",
            "Epoch 17/75\n",
            "49/49 [==============================] - 5s 98ms/step - loss: 1.9068 - accuracy: 0.2984 - val_loss: 1.8549 - val_accuracy: 0.3089\n",
            "Epoch 18/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.8957 - accuracy: 0.3010 - val_loss: 1.8422 - val_accuracy: 0.3157\n",
            "Epoch 19/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.8838 - accuracy: 0.3070 - val_loss: 1.8279 - val_accuracy: 0.3194\n",
            "Epoch 20/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.8722 - accuracy: 0.3101 - val_loss: 1.8166 - val_accuracy: 0.3267\n",
            "Epoch 21/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.8643 - accuracy: 0.3119 - val_loss: 1.8101 - val_accuracy: 0.3241\n",
            "Epoch 22/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.8563 - accuracy: 0.3157 - val_loss: 1.7982 - val_accuracy: 0.3306\n",
            "Epoch 23/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.8421 - accuracy: 0.3237 - val_loss: 1.7873 - val_accuracy: 0.3357\n",
            "Epoch 24/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.8337 - accuracy: 0.3206 - val_loss: 1.7770 - val_accuracy: 0.3415\n",
            "Epoch 25/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.8248 - accuracy: 0.3284 - val_loss: 1.7650 - val_accuracy: 0.3430\n",
            "Epoch 26/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.8179 - accuracy: 0.3280 - val_loss: 1.7575 - val_accuracy: 0.3470\n",
            "Epoch 27/75\n",
            "49/49 [==============================] - 5s 96ms/step - loss: 1.8094 - accuracy: 0.3335 - val_loss: 1.7453 - val_accuracy: 0.3524\n",
            "Epoch 28/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.8007 - accuracy: 0.3351 - val_loss: 1.7403 - val_accuracy: 0.3548\n",
            "Epoch 29/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.7935 - accuracy: 0.3380 - val_loss: 1.7299 - val_accuracy: 0.3583\n",
            "Epoch 30/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.7840 - accuracy: 0.3404 - val_loss: 1.7230 - val_accuracy: 0.3589\n",
            "Epoch 31/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.7748 - accuracy: 0.3472 - val_loss: 1.7088 - val_accuracy: 0.3669\n",
            "Epoch 32/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.7699 - accuracy: 0.3484 - val_loss: 1.7058 - val_accuracy: 0.3662\n",
            "Epoch 33/75\n",
            "49/49 [==============================] - 4s 92ms/step - loss: 1.7625 - accuracy: 0.3499 - val_loss: 1.6921 - val_accuracy: 0.3760\n",
            "Epoch 34/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.7549 - accuracy: 0.3535 - val_loss: 1.6855 - val_accuracy: 0.3765\n",
            "Epoch 35/75\n",
            "49/49 [==============================] - 5s 96ms/step - loss: 1.7466 - accuracy: 0.3550 - val_loss: 1.6798 - val_accuracy: 0.3803\n",
            "Epoch 36/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.7458 - accuracy: 0.3535 - val_loss: 1.6729 - val_accuracy: 0.3832\n",
            "Epoch 37/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.7380 - accuracy: 0.3583 - val_loss: 1.6686 - val_accuracy: 0.3829\n",
            "Epoch 38/75\n",
            "49/49 [==============================] - 5s 105ms/step - loss: 1.7284 - accuracy: 0.3618 - val_loss: 1.6595 - val_accuracy: 0.3879\n",
            "Epoch 39/75\n",
            "49/49 [==============================] - 5s 97ms/step - loss: 1.7246 - accuracy: 0.3632 - val_loss: 1.6531 - val_accuracy: 0.3928\n",
            "Epoch 40/75\n",
            "49/49 [==============================] - 5s 96ms/step - loss: 1.7154 - accuracy: 0.3657 - val_loss: 1.6429 - val_accuracy: 0.3942\n",
            "Epoch 41/75\n",
            "49/49 [==============================] - 5s 98ms/step - loss: 1.7085 - accuracy: 0.3691 - val_loss: 1.6374 - val_accuracy: 0.3969\n",
            "Epoch 42/75\n",
            "49/49 [==============================] - 5s 97ms/step - loss: 1.7036 - accuracy: 0.3705 - val_loss: 1.6336 - val_accuracy: 0.3996\n",
            "Epoch 43/75\n",
            "49/49 [==============================] - 5s 98ms/step - loss: 1.6984 - accuracy: 0.3725 - val_loss: 1.6254 - val_accuracy: 0.4031\n",
            "Epoch 44/75\n",
            "49/49 [==============================] - 5s 96ms/step - loss: 1.6949 - accuracy: 0.3720 - val_loss: 1.6157 - val_accuracy: 0.4057\n",
            "Epoch 45/75\n",
            "49/49 [==============================] - 5s 97ms/step - loss: 1.6887 - accuracy: 0.3742 - val_loss: 1.6113 - val_accuracy: 0.4041\n",
            "Epoch 46/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.6801 - accuracy: 0.3783 - val_loss: 1.6052 - val_accuracy: 0.4063\n",
            "Epoch 47/75\n",
            "49/49 [==============================] - 7s 149ms/step - loss: 1.6769 - accuracy: 0.3811 - val_loss: 1.6038 - val_accuracy: 0.4074\n",
            "Epoch 48/75\n",
            "49/49 [==============================] - 5s 98ms/step - loss: 1.6730 - accuracy: 0.3823 - val_loss: 1.5947 - val_accuracy: 0.4098\n",
            "Epoch 49/75\n",
            "49/49 [==============================] - 4s 83ms/step - loss: 1.6692 - accuracy: 0.3845 - val_loss: 1.5971 - val_accuracy: 0.4115\n",
            "Epoch 50/75\n",
            "49/49 [==============================] - 5s 96ms/step - loss: 1.6612 - accuracy: 0.3849 - val_loss: 1.5874 - val_accuracy: 0.4152\n",
            "Epoch 51/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.6563 - accuracy: 0.3887 - val_loss: 1.5807 - val_accuracy: 0.4159\n",
            "Epoch 52/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.6518 - accuracy: 0.3920 - val_loss: 1.5755 - val_accuracy: 0.4169\n",
            "Epoch 53/75\n",
            "49/49 [==============================] - 5s 96ms/step - loss: 1.6444 - accuracy: 0.3938 - val_loss: 1.5663 - val_accuracy: 0.4205\n",
            "Epoch 54/75\n",
            "49/49 [==============================] - 4s 84ms/step - loss: 1.6444 - accuracy: 0.3917 - val_loss: 1.5692 - val_accuracy: 0.4161\n",
            "Epoch 55/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.6415 - accuracy: 0.3923 - val_loss: 1.5587 - val_accuracy: 0.4233\n",
            "Epoch 56/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.6343 - accuracy: 0.4014 - val_loss: 1.5498 - val_accuracy: 0.4253\n",
            "Epoch 57/75\n",
            "49/49 [==============================] - 4s 82ms/step - loss: 1.6282 - accuracy: 0.4002 - val_loss: 1.5526 - val_accuracy: 0.4271\n",
            "Epoch 58/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.6241 - accuracy: 0.3987 - val_loss: 1.5454 - val_accuracy: 0.4291\n",
            "Epoch 59/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.6193 - accuracy: 0.4045 - val_loss: 1.5336 - val_accuracy: 0.4299\n",
            "Epoch 60/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.6159 - accuracy: 0.4036 - val_loss: 1.5284 - val_accuracy: 0.4358\n",
            "Epoch 61/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.6119 - accuracy: 0.4069 - val_loss: 1.5259 - val_accuracy: 0.4337\n",
            "Epoch 62/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.6055 - accuracy: 0.4109 - val_loss: 1.5237 - val_accuracy: 0.4347\n",
            "Epoch 63/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.6033 - accuracy: 0.4069 - val_loss: 1.5123 - val_accuracy: 0.4405\n",
            "Epoch 64/75\n",
            "49/49 [==============================] - 4s 83ms/step - loss: 1.6002 - accuracy: 0.4082 - val_loss: 1.5283 - val_accuracy: 0.4319\n",
            "Epoch 65/75\n",
            "49/49 [==============================] - 5s 97ms/step - loss: 1.5911 - accuracy: 0.4129 - val_loss: 1.5007 - val_accuracy: 0.4458\n",
            "Epoch 66/75\n",
            "49/49 [==============================] - 4s 84ms/step - loss: 1.5880 - accuracy: 0.4143 - val_loss: 1.5153 - val_accuracy: 0.4390\n",
            "Epoch 67/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.5857 - accuracy: 0.4182 - val_loss: 1.4985 - val_accuracy: 0.4427\n",
            "Epoch 68/75\n",
            "49/49 [==============================] - 4s 82ms/step - loss: 1.5832 - accuracy: 0.4180 - val_loss: 1.5257 - val_accuracy: 0.4387\n",
            "Epoch 69/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.5799 - accuracy: 0.4214 - val_loss: 1.4950 - val_accuracy: 0.4474\n",
            "Epoch 70/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.5709 - accuracy: 0.4238 - val_loss: 1.4795 - val_accuracy: 0.4507\n",
            "Epoch 71/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 1.5674 - accuracy: 0.4223 - val_loss: 1.4776 - val_accuracy: 0.4519\n",
            "Epoch 72/75\n",
            "49/49 [==============================] - 4s 84ms/step - loss: 1.5678 - accuracy: 0.4222 - val_loss: 1.4874 - val_accuracy: 0.4478\n",
            "Epoch 73/75\n",
            "49/49 [==============================] - 4s 82ms/step - loss: 1.5613 - accuracy: 0.4254 - val_loss: 1.5141 - val_accuracy: 0.4521\n",
            "Epoch 74/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 1.5572 - accuracy: 0.4259 - val_loss: 1.4686 - val_accuracy: 0.4565\n",
            "Epoch 75/75\n",
            "49/49 [==============================] - 4s 83ms/step - loss: 1.5508 - accuracy: 0.4282 - val_loss: 1.4838 - val_accuracy: 0.4500\n",
            "Network takes 373.344 seconds to train\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2GI0ADGIiil1",
        "colab_type": "code",
        "outputId": "efa39f1a-08d5-4c66-b91e-184614b144d0",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 68
        }
      },
      "source": [
        "model.save_weights(\"resnet_lrsch_2.h5\")\n",
        "!gsutil -m cp resnet_lrsch_2.h5 gs://batch_norm_tpu/"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Copying file://resnet_lrsch_2.h5 [Content-Type=application/octet-stream]...\n",
            "/ [1/1 files][  2.4 MiB/  2.4 MiB] 100% Done                                    \n",
            "Operation completed over 1 objects/2.4 MiB.                                      \n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ou8GN5VhWrbd",
        "colab_type": "code",
        "outputId": "4492ed98-144a-4dbd-fad9-33d06d45a4ed",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        }
      },
      "source": [
        "with strategy.scope():\n",
        "    model = get_training_model()\n",
        "\n",
        "    for layer in model.layers:\n",
        "        if not isinstance(layer, tf.keras.layers.BatchNormalization):\n",
        "            if hasattr(layer, \"trainable\"):\n",
        "                layer.trainable=False\n",
        "    model.compile(loss=\"sparse_categorical_crossentropy\", optimizer=\"sgd\", metrics=[\"accuracy\"])\n",
        "    \n",
        "model.summary()"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Model: \"model_7\"\n",
            "__________________________________________________________________________________________________\n",
            "Layer (type)                    Output Shape         Param #     Connected to                     \n",
            "==================================================================================================\n",
            "input_8 (InputLayer)            [(None, 32, 32, 3)]  0                                            \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_298 (Conv2D)             (None, 32, 32, 16)   448         input_8[0][0]                    \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_284 (BatchN (None, 32, 32, 16)   64          conv2d_298[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_284 (ReLU)                (None, 32, 32, 16)   0           batch_normalization_284[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_300 (Conv2D)             (None, 32, 32, 16)   272         re_lu_284[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_285 (BatchN (None, 32, 32, 16)   64          conv2d_300[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_285 (ReLU)                (None, 32, 32, 16)   0           batch_normalization_285[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_301 (Conv2D)             (None, 32, 32, 16)   2320        re_lu_285[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_286 (BatchN (None, 32, 32, 16)   64          conv2d_301[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_286 (ReLU)                (None, 32, 32, 16)   0           batch_normalization_286[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_302 (Conv2D)             (None, 32, 32, 64)   1088        re_lu_286[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_299 (Conv2D)             (None, 32, 32, 64)   1088        re_lu_284[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_287 (BatchN (None, 32, 32, 64)   256         conv2d_302[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "add_90 (Add)                    (None, 32, 32, 64)   0           conv2d_299[0][0]                 \n",
            "                                                                 batch_normalization_287[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_287 (ReLU)                (None, 32, 32, 64)   0           add_90[0][0]                     \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_303 (Conv2D)             (None, 32, 32, 16)   1040        re_lu_287[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_288 (BatchN (None, 32, 32, 16)   64          conv2d_303[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_288 (ReLU)                (None, 32, 32, 16)   0           batch_normalization_288[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_304 (Conv2D)             (None, 32, 32, 16)   2320        re_lu_288[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_289 (BatchN (None, 32, 32, 16)   64          conv2d_304[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_289 (ReLU)                (None, 32, 32, 16)   0           batch_normalization_289[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_305 (Conv2D)             (None, 32, 32, 64)   1088        re_lu_289[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_290 (BatchN (None, 32, 32, 64)   256         conv2d_305[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "add_91 (Add)                    (None, 32, 32, 64)   0           batch_normalization_290[0][0]    \n",
            "                                                                 re_lu_287[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_290 (ReLU)                (None, 32, 32, 64)   0           add_91[0][0]                     \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_307 (Conv2D)             (None, 32, 32, 64)   4160        re_lu_290[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_291 (BatchN (None, 32, 32, 64)   256         conv2d_307[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_291 (ReLU)                (None, 32, 32, 64)   0           batch_normalization_291[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_308 (Conv2D)             (None, 16, 16, 64)   36928       re_lu_291[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_292 (BatchN (None, 16, 16, 64)   256         conv2d_308[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_292 (ReLU)                (None, 16, 16, 64)   0           batch_normalization_292[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_309 (Conv2D)             (None, 16, 16, 128)  8320        re_lu_292[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_306 (Conv2D)             (None, 16, 16, 128)  8320        re_lu_290[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_293 (BatchN (None, 16, 16, 128)  512         conv2d_309[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "add_92 (Add)                    (None, 16, 16, 128)  0           conv2d_306[0][0]                 \n",
            "                                                                 batch_normalization_293[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_293 (ReLU)                (None, 16, 16, 128)  0           add_92[0][0]                     \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_310 (Conv2D)             (None, 16, 16, 64)   8256        re_lu_293[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_294 (BatchN (None, 16, 16, 64)   256         conv2d_310[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_294 (ReLU)                (None, 16, 16, 64)   0           batch_normalization_294[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_311 (Conv2D)             (None, 16, 16, 64)   36928       re_lu_294[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_295 (BatchN (None, 16, 16, 64)   256         conv2d_311[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_295 (ReLU)                (None, 16, 16, 64)   0           batch_normalization_295[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_312 (Conv2D)             (None, 16, 16, 128)  8320        re_lu_295[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_296 (BatchN (None, 16, 16, 128)  512         conv2d_312[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "add_93 (Add)                    (None, 16, 16, 128)  0           batch_normalization_296[0][0]    \n",
            "                                                                 re_lu_293[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_296 (ReLU)                (None, 16, 16, 128)  0           add_93[0][0]                     \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_314 (Conv2D)             (None, 16, 16, 128)  16512       re_lu_296[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_297 (BatchN (None, 16, 16, 128)  512         conv2d_314[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_297 (ReLU)                (None, 16, 16, 128)  0           batch_normalization_297[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_315 (Conv2D)             (None, 8, 8, 128)    147584      re_lu_297[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_298 (BatchN (None, 8, 8, 128)    512         conv2d_315[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_298 (ReLU)                (None, 8, 8, 128)    0           batch_normalization_298[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_316 (Conv2D)             (None, 8, 8, 256)    33024       re_lu_298[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_313 (Conv2D)             (None, 8, 8, 256)    33024       re_lu_296[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_299 (BatchN (None, 8, 8, 256)    1024        conv2d_316[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "add_94 (Add)                    (None, 8, 8, 256)    0           conv2d_313[0][0]                 \n",
            "                                                                 batch_normalization_299[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_299 (ReLU)                (None, 8, 8, 256)    0           add_94[0][0]                     \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_317 (Conv2D)             (None, 8, 8, 128)    32896       re_lu_299[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_300 (BatchN (None, 8, 8, 128)    512         conv2d_317[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_300 (ReLU)                (None, 8, 8, 128)    0           batch_normalization_300[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_318 (Conv2D)             (None, 8, 8, 128)    147584      re_lu_300[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_301 (BatchN (None, 8, 8, 128)    512         conv2d_318[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_301 (ReLU)                (None, 8, 8, 128)    0           batch_normalization_301[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_319 (Conv2D)             (None, 8, 8, 256)    33024       re_lu_301[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_302 (BatchN (None, 8, 8, 256)    1024        conv2d_319[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "add_95 (Add)                    (None, 8, 8, 256)    0           batch_normalization_302[0][0]    \n",
            "                                                                 re_lu_299[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_302 (ReLU)                (None, 8, 8, 256)    0           add_95[0][0]                     \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_303 (BatchN (None, 8, 8, 256)    1024        re_lu_302[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "re_lu_303 (ReLU)                (None, 8, 8, 256)    0           batch_normalization_303[0][0]    \n",
            "__________________________________________________________________________________________________\n",
            "average_pooling2d_7 (AveragePoo (None, 1, 1, 256)    0           re_lu_303[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "flatten_7 (Flatten)             (None, 256)          0           average_pooling2d_7[0][0]        \n",
            "__________________________________________________________________________________________________\n",
            "dense_7 (Dense)                 (None, 10)           2570        flatten_7[0][0]                  \n",
            "==================================================================================================\n",
            "Total params: 575,114\n",
            "Trainable params: 4,000\n",
            "Non-trainable params: 571,114\n",
            "__________________________________________________________________________________________________\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q0pgz6fvajZV",
        "colab_type": "text"
      },
      "source": [
        "- Total params: 575,114\n",
        "- **Trainable params: 4,000**\n",
        "- Non-trainable params: 571,114"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HpzMXab8Us7C",
        "colab_type": "code",
        "outputId": "fbdb6130-f817-4ad0-be92-45251e4d97f4",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        }
      },
      "source": [
        "wandb.init(project=\"training-bn-only\", id=\"bn-only-vanilla-1\")\n",
        "\n",
        "start = time.time()\n",
        "h = model.fit(train_ds,\n",
        "         validation_data=test_ds,\n",
        "         epochs=75,\n",
        "         callbacks=[WandbCallback()])\n",
        "end = time.time()\n",
        "wandb.log({\"training_time\": end - start})\n",
        "print(\"Network takes {:.3f} seconds to train\".format(end - start))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "                Logging results to <a href=\"https://wandb.com\" target=\"_blank\">Weights & Biases</a> <a href=\"https://docs.wandb.com/integrations/jupyter.html\" target=\"_blank\">(Documentation)</a>.<br/>\n",
              "                Project page: <a href=\"https://app.wandb.ai/sayakpaul/training-bn-only\" target=\"_blank\">https://app.wandb.ai/sayakpaul/training-bn-only</a><br/>\n",
              "                Run page: <a href=\"https://app.wandb.ai/sayakpaul/training-bn-only/runs/bn-only-vanilla-1\" target=\"_blank\">https://app.wandb.ai/sayakpaul/training-bn-only/runs/bn-only-vanilla-1</a><br/>\n",
              "            "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "Epoch 1/75\n",
            "49/49 [==============================] - 9s 189ms/step - loss: 2.4312 - accuracy: 0.0752 - val_loss: 2.5910 - val_accuracy: 0.0929\n",
            "Epoch 2/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 2.4248 - accuracy: 0.0747 - val_loss: 2.3635 - val_accuracy: 0.0990\n",
            "Epoch 3/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 2.4197 - accuracy: 0.0755 - val_loss: 2.3360 - val_accuracy: 0.0988\n",
            "Epoch 4/75\n",
            "49/49 [==============================] - 4s 81ms/step - loss: 2.4156 - accuracy: 0.0763 - val_loss: 2.3367 - val_accuracy: 0.0828\n",
            "Epoch 5/75\n",
            "49/49 [==============================] - 4s 79ms/step - loss: 2.4099 - accuracy: 0.0777 - val_loss: 2.3374 - val_accuracy: 0.0769\n",
            "Epoch 6/75\n",
            "49/49 [==============================] - 4s 79ms/step - loss: 2.4082 - accuracy: 0.0787 - val_loss: 2.3373 - val_accuracy: 0.0776\n",
            "Epoch 7/75\n",
            "49/49 [==============================] - 4s 78ms/step - loss: 2.3998 - accuracy: 0.0796 - val_loss: 2.3378 - val_accuracy: 0.0825\n",
            "Epoch 8/75\n",
            "49/49 [==============================] - 4s 79ms/step - loss: 2.3944 - accuracy: 0.0829 - val_loss: 2.3390 - val_accuracy: 0.0845\n",
            "Epoch 9/75\n",
            "49/49 [==============================] - 4s 81ms/step - loss: 2.3945 - accuracy: 0.0806 - val_loss: 2.3392 - val_accuracy: 0.0879\n",
            "Epoch 10/75\n",
            "49/49 [==============================] - 4s 79ms/step - loss: 2.3897 - accuracy: 0.0833 - val_loss: 2.3390 - val_accuracy: 0.0920\n",
            "Epoch 11/75\n",
            "49/49 [==============================] - 4s 79ms/step - loss: 2.3866 - accuracy: 0.0838 - val_loss: 2.3382 - val_accuracy: 0.0936\n",
            "Epoch 12/75\n",
            "49/49 [==============================] - 4s 79ms/step - loss: 2.3820 - accuracy: 0.0885 - val_loss: 2.3365 - val_accuracy: 0.0980\n",
            "Epoch 13/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 2.3792 - accuracy: 0.0862 - val_loss: 2.3347 - val_accuracy: 0.0994\n",
            "Epoch 14/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 2.3731 - accuracy: 0.0902 - val_loss: 2.3326 - val_accuracy: 0.1022\n",
            "Epoch 15/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.3702 - accuracy: 0.0921 - val_loss: 2.3303 - val_accuracy: 0.1047\n",
            "Epoch 16/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 2.3670 - accuracy: 0.0939 - val_loss: 2.3282 - val_accuracy: 0.1057\n",
            "Epoch 17/75\n",
            "49/49 [==============================] - 4s 90ms/step - loss: 2.3653 - accuracy: 0.0939 - val_loss: 2.3261 - val_accuracy: 0.1087\n",
            "Epoch 18/75\n",
            "49/49 [==============================] - 5s 92ms/step - loss: 2.3606 - accuracy: 0.0980 - val_loss: 2.3238 - val_accuracy: 0.1108\n",
            "Epoch 19/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.3582 - accuracy: 0.0962 - val_loss: 2.3218 - val_accuracy: 0.1125\n",
            "Epoch 20/75\n",
            "49/49 [==============================] - 5s 92ms/step - loss: 2.3542 - accuracy: 0.0998 - val_loss: 2.3197 - val_accuracy: 0.1151\n",
            "Epoch 21/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.3529 - accuracy: 0.0995 - val_loss: 2.3176 - val_accuracy: 0.1174\n",
            "Epoch 22/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.3513 - accuracy: 0.1015 - val_loss: 2.3156 - val_accuracy: 0.1208\n",
            "Epoch 23/75\n",
            "49/49 [==============================] - 4s 92ms/step - loss: 2.3484 - accuracy: 0.1032 - val_loss: 2.3136 - val_accuracy: 0.1222\n",
            "Epoch 24/75\n",
            "49/49 [==============================] - 4s 90ms/step - loss: 2.3480 - accuracy: 0.1042 - val_loss: 2.3116 - val_accuracy: 0.1234\n",
            "Epoch 25/75\n",
            "49/49 [==============================] - 9s 183ms/step - loss: 2.3434 - accuracy: 0.1078 - val_loss: 2.3099 - val_accuracy: 0.1257\n",
            "Epoch 26/75\n",
            "49/49 [==============================] - 5s 92ms/step - loss: 2.3414 - accuracy: 0.1080 - val_loss: 2.3082 - val_accuracy: 0.1271\n",
            "Epoch 27/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 2.3382 - accuracy: 0.1106 - val_loss: 2.3065 - val_accuracy: 0.1296\n",
            "Epoch 28/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.3349 - accuracy: 0.1129 - val_loss: 2.3049 - val_accuracy: 0.1320\n",
            "Epoch 29/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.3335 - accuracy: 0.1152 - val_loss: 2.3032 - val_accuracy: 0.1330\n",
            "Epoch 30/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 2.3343 - accuracy: 0.1169 - val_loss: 2.3016 - val_accuracy: 0.1366\n",
            "Epoch 31/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.3304 - accuracy: 0.1171 - val_loss: 2.3002 - val_accuracy: 0.1384\n",
            "Epoch 32/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.3279 - accuracy: 0.1182 - val_loss: 2.2986 - val_accuracy: 0.1401\n",
            "Epoch 33/75\n",
            "49/49 [==============================] - 4s 92ms/step - loss: 2.3262 - accuracy: 0.1191 - val_loss: 2.2969 - val_accuracy: 0.1450\n",
            "Epoch 34/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.3252 - accuracy: 0.1186 - val_loss: 2.2955 - val_accuracy: 0.1453\n",
            "Epoch 35/75\n",
            "49/49 [==============================] - 4s 92ms/step - loss: 2.3243 - accuracy: 0.1176 - val_loss: 2.2939 - val_accuracy: 0.1487\n",
            "Epoch 36/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.3217 - accuracy: 0.1246 - val_loss: 2.2926 - val_accuracy: 0.1496\n",
            "Epoch 37/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 2.3180 - accuracy: 0.1275 - val_loss: 2.2914 - val_accuracy: 0.1507\n",
            "Epoch 38/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 2.3180 - accuracy: 0.1284 - val_loss: 2.2898 - val_accuracy: 0.1523\n",
            "Epoch 39/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 2.3154 - accuracy: 0.1292 - val_loss: 2.2886 - val_accuracy: 0.1540\n",
            "Epoch 40/75\n",
            "49/49 [==============================] - 5s 95ms/step - loss: 2.3158 - accuracy: 0.1270 - val_loss: 2.2873 - val_accuracy: 0.1555\n",
            "Epoch 41/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.3120 - accuracy: 0.1310 - val_loss: 2.2861 - val_accuracy: 0.1565\n",
            "Epoch 42/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.3103 - accuracy: 0.1332 - val_loss: 2.2849 - val_accuracy: 0.1567\n",
            "Epoch 43/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.3105 - accuracy: 0.1334 - val_loss: 2.2837 - val_accuracy: 0.1593\n",
            "Epoch 44/75\n",
            "49/49 [==============================] - 5s 92ms/step - loss: 2.3076 - accuracy: 0.1335 - val_loss: 2.2825 - val_accuracy: 0.1597\n",
            "Epoch 45/75\n",
            "49/49 [==============================] - 4s 90ms/step - loss: 2.3083 - accuracy: 0.1326 - val_loss: 2.2812 - val_accuracy: 0.1626\n",
            "Epoch 46/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.3066 - accuracy: 0.1361 - val_loss: 2.2800 - val_accuracy: 0.1624\n",
            "Epoch 47/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.3048 - accuracy: 0.1366 - val_loss: 2.2790 - val_accuracy: 0.1624\n",
            "Epoch 48/75\n",
            "49/49 [==============================] - 4s 90ms/step - loss: 2.3023 - accuracy: 0.1376 - val_loss: 2.2779 - val_accuracy: 0.1622\n",
            "Epoch 49/75\n",
            "49/49 [==============================] - 5s 92ms/step - loss: 2.3014 - accuracy: 0.1394 - val_loss: 2.2769 - val_accuracy: 0.1644\n",
            "Epoch 50/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.2997 - accuracy: 0.1381 - val_loss: 2.2757 - val_accuracy: 0.1648\n",
            "Epoch 51/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.2978 - accuracy: 0.1412 - val_loss: 2.2745 - val_accuracy: 0.1644\n",
            "Epoch 52/75\n",
            "49/49 [==============================] - 5s 97ms/step - loss: 2.2976 - accuracy: 0.1393 - val_loss: 2.2734 - val_accuracy: 0.1650\n",
            "Epoch 53/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 2.2954 - accuracy: 0.1440 - val_loss: 2.2724 - val_accuracy: 0.1662\n",
            "Epoch 54/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 2.2943 - accuracy: 0.1434 - val_loss: 2.2714 - val_accuracy: 0.1682\n",
            "Epoch 55/75\n",
            "49/49 [==============================] - 4s 90ms/step - loss: 2.2930 - accuracy: 0.1444 - val_loss: 2.2701 - val_accuracy: 0.1674\n",
            "Epoch 56/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.2912 - accuracy: 0.1460 - val_loss: 2.2692 - val_accuracy: 0.1692\n",
            "Epoch 57/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.2918 - accuracy: 0.1450 - val_loss: 2.2680 - val_accuracy: 0.1696\n",
            "Epoch 58/75\n",
            "49/49 [==============================] - 4s 90ms/step - loss: 2.2901 - accuracy: 0.1484 - val_loss: 2.2670 - val_accuracy: 0.1714\n",
            "Epoch 59/75\n",
            "49/49 [==============================] - 4s 89ms/step - loss: 2.2892 - accuracy: 0.1467 - val_loss: 2.2661 - val_accuracy: 0.1725\n",
            "Epoch 60/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.2885 - accuracy: 0.1468 - val_loss: 2.2652 - val_accuracy: 0.1720\n",
            "Epoch 61/75\n",
            "49/49 [==============================] - 5s 92ms/step - loss: 2.2870 - accuracy: 0.1487 - val_loss: 2.2643 - val_accuracy: 0.1736\n",
            "Epoch 62/75\n",
            "49/49 [==============================] - 5s 92ms/step - loss: 2.2856 - accuracy: 0.1503 - val_loss: 2.2634 - val_accuracy: 0.1728\n",
            "Epoch 63/75\n",
            "49/49 [==============================] - 4s 90ms/step - loss: 2.2860 - accuracy: 0.1504 - val_loss: 2.2623 - val_accuracy: 0.1721\n",
            "Epoch 64/75\n",
            "49/49 [==============================] - 4s 90ms/step - loss: 2.2816 - accuracy: 0.1535 - val_loss: 2.2614 - val_accuracy: 0.1733\n",
            "Epoch 65/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.2845 - accuracy: 0.1502 - val_loss: 2.2607 - val_accuracy: 0.1751\n",
            "Epoch 66/75\n",
            "49/49 [==============================] - 5s 92ms/step - loss: 2.2815 - accuracy: 0.1505 - val_loss: 2.2597 - val_accuracy: 0.1757\n",
            "Epoch 67/75\n",
            "49/49 [==============================] - 4s 89ms/step - loss: 2.2799 - accuracy: 0.1535 - val_loss: 2.2588 - val_accuracy: 0.1764\n",
            "Epoch 68/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.2796 - accuracy: 0.1506 - val_loss: 2.2578 - val_accuracy: 0.1773\n",
            "Epoch 69/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.2796 - accuracy: 0.1532 - val_loss: 2.2571 - val_accuracy: 0.1775\n",
            "Epoch 70/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.2775 - accuracy: 0.1537 - val_loss: 2.2563 - val_accuracy: 0.1779\n",
            "Epoch 71/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 2.2769 - accuracy: 0.1553 - val_loss: 2.2553 - val_accuracy: 0.1788\n",
            "Epoch 72/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 2.2750 - accuracy: 0.1560 - val_loss: 2.2546 - val_accuracy: 0.1788\n",
            "Epoch 73/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 2.2751 - accuracy: 0.1548 - val_loss: 2.2537 - val_accuracy: 0.1789\n",
            "Epoch 74/75\n",
            "49/49 [==============================] - 4s 90ms/step - loss: 2.2757 - accuracy: 0.1540 - val_loss: 2.2530 - val_accuracy: 0.1780\n",
            "Epoch 75/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.2728 - accuracy: 0.1584 - val_loss: 2.2522 - val_accuracy: 0.1802\n",
            "Network takes 356.103 seconds to train\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2PXol7EltH0s",
        "colab_type": "code",
        "outputId": "b1986cc6-46c9-4a6b-91d6-b9bb963b9068",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 68
        }
      },
      "source": [
        "model.save_weights(\"resnet_bn_only.h5\")\n",
        "!gsutil -m cp -r resnet_bn_only.h5 gs://batch_norm_tpu/"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Copying file://resnet_bn_only.h5 [Content-Type=application/octet-stream]...\n",
            "/ [1/1 files][  2.4 MiB/  2.4 MiB] 100% Done                                    \n",
            "Operation completed over 1 objects/2.4 MiB.                                      \n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "D-kcR7X1b3CO",
        "colab_type": "code",
        "outputId": "ff4701eb-a5a2-4953-a6b3-19ece6788288",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        }
      },
      "source": [
        "# Train model with a decay schedule\n",
        "wandb.init(project=\"training-bn-only\", id=\"bn-only-lrsch-1\")\n",
        "\n",
        "first_decay_steps = 1000\n",
        "lr_decayed_fn = (\n",
        "  tf.keras.experimental.CosineDecayRestarts(\n",
        "      initial_learning_rate=0.1,\n",
        "      first_decay_steps=first_decay_steps))\n",
        "\n",
        "with strategy.scope():\n",
        "    model = get_training_model()\n",
        "\n",
        "    for layer in model.layers:\n",
        "        if not isinstance(layer, tf.keras.layers.BatchNormalization):\n",
        "            if hasattr(layer, \"trainable\"):\n",
        "                layer.trainable=False\n",
        "    model.compile(loss=\"sparse_categorical_crossentropy\", optimizer=tf.keras.optimizers.SGD(lr_decayed_fn), metrics=[\"accuracy\"])\n",
        "    \n",
        "start = time.time()\n",
        "h = model.fit(train_ds,\n",
        "         validation_data=test_ds,\n",
        "         epochs=75,\n",
        "         callbacks=[WandbCallback()])\n",
        "end = time.time()\n",
        "wandb.log({\"training_time\": end - start})\n",
        "print(\"Network takes {:.3f} seconds to train\".format(end - start))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "                Logging results to <a href=\"https://wandb.com\" target=\"_blank\">Weights & Biases</a> <a href=\"https://docs.wandb.com/integrations/jupyter.html\" target=\"_blank\">(Documentation)</a>.<br/>\n",
              "                Project page: <a href=\"https://app.wandb.ai/sayakpaul/training-bn-only\" target=\"_blank\">https://app.wandb.ai/sayakpaul/training-bn-only</a><br/>\n",
              "                Run page: <a href=\"https://app.wandb.ai/sayakpaul/training-bn-only/runs/bn-only-lrsch-1\" target=\"_blank\">https://app.wandb.ai/sayakpaul/training-bn-only/runs/bn-only-lrsch-1</a><br/>\n",
              "            "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "Epoch 1/75\n",
            "49/49 [==============================] - 10s 200ms/step - loss: 2.4715 - accuracy: 0.0960 - val_loss: 2.3537 - val_accuracy: 0.1055\n",
            "Epoch 2/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.3948 - accuracy: 0.1009 - val_loss: 2.3078 - val_accuracy: 0.1105\n",
            "Epoch 3/75\n",
            "49/49 [==============================] - 5s 92ms/step - loss: 2.3469 - accuracy: 0.1148 - val_loss: 2.2999 - val_accuracy: 0.1224\n",
            "Epoch 4/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.3161 - accuracy: 0.1236 - val_loss: 2.2919 - val_accuracy: 0.1286\n",
            "Epoch 5/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 2.2959 - accuracy: 0.1358 - val_loss: 2.2822 - val_accuracy: 0.1375\n",
            "Epoch 6/75\n",
            "49/49 [==============================] - 4s 92ms/step - loss: 2.2790 - accuracy: 0.1405 - val_loss: 2.2685 - val_accuracy: 0.1497\n",
            "Epoch 7/75\n",
            "49/49 [==============================] - 5s 92ms/step - loss: 2.2676 - accuracy: 0.1463 - val_loss: 2.2553 - val_accuracy: 0.1584\n",
            "Epoch 8/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 2.2591 - accuracy: 0.1507 - val_loss: 2.2448 - val_accuracy: 0.1640\n",
            "Epoch 9/75\n",
            "49/49 [==============================] - 5s 92ms/step - loss: 2.2523 - accuracy: 0.1528 - val_loss: 2.2355 - val_accuracy: 0.1708\n",
            "Epoch 10/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.2454 - accuracy: 0.1599 - val_loss: 2.2286 - val_accuracy: 0.1730\n",
            "Epoch 11/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 2.2398 - accuracy: 0.1623 - val_loss: 2.2224 - val_accuracy: 0.1780\n",
            "Epoch 12/75\n",
            "49/49 [==============================] - 5s 92ms/step - loss: 2.2366 - accuracy: 0.1680 - val_loss: 2.2179 - val_accuracy: 0.1835\n",
            "Epoch 13/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.2341 - accuracy: 0.1668 - val_loss: 2.2141 - val_accuracy: 0.1828\n",
            "Epoch 14/75\n",
            "49/49 [==============================] - 4s 92ms/step - loss: 2.2300 - accuracy: 0.1689 - val_loss: 2.2108 - val_accuracy: 0.1852\n",
            "Epoch 15/75\n",
            "49/49 [==============================] - 4s 90ms/step - loss: 2.2282 - accuracy: 0.1692 - val_loss: 2.2087 - val_accuracy: 0.1871\n",
            "Epoch 16/75\n",
            "49/49 [==============================] - 10s 199ms/step - loss: 2.2256 - accuracy: 0.1720 - val_loss: 2.2077 - val_accuracy: 0.1874\n",
            "Epoch 17/75\n",
            "49/49 [==============================] - 5s 92ms/step - loss: 2.2253 - accuracy: 0.1699 - val_loss: 2.2069 - val_accuracy: 0.1866\n",
            "Epoch 18/75\n",
            "49/49 [==============================] - 4s 92ms/step - loss: 2.2231 - accuracy: 0.1730 - val_loss: 2.2063 - val_accuracy: 0.1868\n",
            "Epoch 19/75\n",
            "49/49 [==============================] - 5s 92ms/step - loss: 2.2238 - accuracy: 0.1726 - val_loss: 2.2062 - val_accuracy: 0.1865\n",
            "Epoch 20/75\n",
            "49/49 [==============================] - 4s 81ms/step - loss: 2.2230 - accuracy: 0.1725 - val_loss: 2.2062 - val_accuracy: 0.1857\n",
            "Epoch 21/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 2.2225 - accuracy: 0.1721 - val_loss: 2.1996 - val_accuracy: 0.1888\n",
            "Epoch 22/75\n",
            "49/49 [==============================] - 5s 92ms/step - loss: 2.2162 - accuracy: 0.1768 - val_loss: 2.1908 - val_accuracy: 0.1905\n",
            "Epoch 23/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 2.2075 - accuracy: 0.1789 - val_loss: 2.1800 - val_accuracy: 0.1922\n",
            "Epoch 24/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 2.1997 - accuracy: 0.1848 - val_loss: 2.1704 - val_accuracy: 0.1969\n",
            "Epoch 25/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.1940 - accuracy: 0.1875 - val_loss: 2.1627 - val_accuracy: 0.1967\n",
            "Epoch 26/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.1864 - accuracy: 0.1902 - val_loss: 2.1551 - val_accuracy: 0.2017\n",
            "Epoch 27/75\n",
            "49/49 [==============================] - 5s 92ms/step - loss: 2.1784 - accuracy: 0.1924 - val_loss: 2.1481 - val_accuracy: 0.2027\n",
            "Epoch 28/75\n",
            "49/49 [==============================] - 4s 90ms/step - loss: 2.1718 - accuracy: 0.1945 - val_loss: 2.1424 - val_accuracy: 0.2035\n",
            "Epoch 29/75\n",
            "49/49 [==============================] - 5s 92ms/step - loss: 2.1680 - accuracy: 0.1939 - val_loss: 2.1356 - val_accuracy: 0.2042\n",
            "Epoch 30/75\n",
            "49/49 [==============================] - 5s 92ms/step - loss: 2.1619 - accuracy: 0.1987 - val_loss: 2.1305 - val_accuracy: 0.2052\n",
            "Epoch 31/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.1586 - accuracy: 0.1996 - val_loss: 2.1257 - val_accuracy: 0.2060\n",
            "Epoch 32/75\n",
            "49/49 [==============================] - 4s 90ms/step - loss: 2.1556 - accuracy: 0.2012 - val_loss: 2.1211 - val_accuracy: 0.2053\n",
            "Epoch 33/75\n",
            "49/49 [==============================] - 4s 92ms/step - loss: 2.1488 - accuracy: 0.2025 - val_loss: 2.1163 - val_accuracy: 0.2074\n",
            "Epoch 34/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.1476 - accuracy: 0.1994 - val_loss: 2.1118 - val_accuracy: 0.2066\n",
            "Epoch 35/75\n",
            "49/49 [==============================] - 4s 90ms/step - loss: 2.1439 - accuracy: 0.2026 - val_loss: 2.1097 - val_accuracy: 0.2086\n",
            "Epoch 36/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.1381 - accuracy: 0.2048 - val_loss: 2.1044 - val_accuracy: 0.2081\n",
            "Epoch 37/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 2.1354 - accuracy: 0.2071 - val_loss: 2.1009 - val_accuracy: 0.2091\n",
            "Epoch 38/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 2.1340 - accuracy: 0.2045 - val_loss: 2.0985 - val_accuracy: 0.2129\n",
            "Epoch 39/75\n",
            "49/49 [==============================] - 4s 92ms/step - loss: 2.1308 - accuracy: 0.2080 - val_loss: 2.0954 - val_accuracy: 0.2131\n",
            "Epoch 40/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 2.1289 - accuracy: 0.2073 - val_loss: 2.0923 - val_accuracy: 0.2157\n",
            "Epoch 41/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 2.1232 - accuracy: 0.2101 - val_loss: 2.0903 - val_accuracy: 0.2147\n",
            "Epoch 42/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 2.1233 - accuracy: 0.2075 - val_loss: 2.0871 - val_accuracy: 0.2175\n",
            "Epoch 43/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.1226 - accuracy: 0.2101 - val_loss: 2.0860 - val_accuracy: 0.2175\n",
            "Epoch 44/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 2.1204 - accuracy: 0.2123 - val_loss: 2.0843 - val_accuracy: 0.2179\n",
            "Epoch 45/75\n",
            "49/49 [==============================] - 4s 92ms/step - loss: 2.1181 - accuracy: 0.2116 - val_loss: 2.0822 - val_accuracy: 0.2195\n",
            "Epoch 46/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.1153 - accuracy: 0.2128 - val_loss: 2.0804 - val_accuracy: 0.2197\n",
            "Epoch 47/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.1142 - accuracy: 0.2156 - val_loss: 2.0787 - val_accuracy: 0.2208\n",
            "Epoch 48/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.1139 - accuracy: 0.2127 - val_loss: 2.0779 - val_accuracy: 0.2206\n",
            "Epoch 49/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.1114 - accuracy: 0.2130 - val_loss: 2.0775 - val_accuracy: 0.2204\n",
            "Epoch 50/75\n",
            "49/49 [==============================] - 4s 90ms/step - loss: 2.1120 - accuracy: 0.2137 - val_loss: 2.0764 - val_accuracy: 0.2199\n",
            "Epoch 51/75\n",
            "49/49 [==============================] - 4s 90ms/step - loss: 2.1116 - accuracy: 0.2152 - val_loss: 2.0755 - val_accuracy: 0.2187\n",
            "Epoch 52/75\n",
            "49/49 [==============================] - 4s 92ms/step - loss: 2.1103 - accuracy: 0.2180 - val_loss: 2.0749 - val_accuracy: 0.2192\n",
            "Epoch 53/75\n",
            "49/49 [==============================] - 4s 90ms/step - loss: 2.1091 - accuracy: 0.2154 - val_loss: 2.0744 - val_accuracy: 0.2193\n",
            "Epoch 54/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.1097 - accuracy: 0.2138 - val_loss: 2.0738 - val_accuracy: 0.2188\n",
            "Epoch 55/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 2.1089 - accuracy: 0.2139 - val_loss: 2.0736 - val_accuracy: 0.2189\n",
            "Epoch 56/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.1096 - accuracy: 0.2139 - val_loss: 2.0734 - val_accuracy: 0.2212\n",
            "Epoch 57/75\n",
            "49/49 [==============================] - 4s 92ms/step - loss: 2.1116 - accuracy: 0.2124 - val_loss: 2.0733 - val_accuracy: 0.2189\n",
            "Epoch 58/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.1119 - accuracy: 0.2126 - val_loss: 2.0733 - val_accuracy: 0.2186\n",
            "Epoch 59/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 2.1070 - accuracy: 0.2152 - val_loss: 2.0732 - val_accuracy: 0.2187\n",
            "Epoch 60/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 2.1084 - accuracy: 0.2136 - val_loss: 2.0730 - val_accuracy: 0.2200\n",
            "Epoch 61/75\n",
            "49/49 [==============================] - 4s 78ms/step - loss: 2.1079 - accuracy: 0.2143 - val_loss: 2.0730 - val_accuracy: 0.2189\n",
            "Epoch 62/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.1060 - accuracy: 0.2171 - val_loss: 2.0688 - val_accuracy: 0.2209\n",
            "Epoch 63/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 2.1041 - accuracy: 0.2142 - val_loss: 2.0645 - val_accuracy: 0.2220\n",
            "Epoch 64/75\n",
            "49/49 [==============================] - 4s 90ms/step - loss: 2.0981 - accuracy: 0.2176 - val_loss: 2.0615 - val_accuracy: 0.2207\n",
            "Epoch 65/75\n",
            "49/49 [==============================] - 5s 94ms/step - loss: 2.0969 - accuracy: 0.2182 - val_loss: 2.0609 - val_accuracy: 0.2242\n",
            "Epoch 66/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 2.0915 - accuracy: 0.2202 - val_loss: 2.0541 - val_accuracy: 0.2275\n",
            "Epoch 67/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.0902 - accuracy: 0.2218 - val_loss: 2.0516 - val_accuracy: 0.2281\n",
            "Epoch 68/75\n",
            "49/49 [==============================] - 4s 90ms/step - loss: 2.0857 - accuracy: 0.2206 - val_loss: 2.0478 - val_accuracy: 0.2276\n",
            "Epoch 69/75\n",
            "49/49 [==============================] - 5s 93ms/step - loss: 2.0794 - accuracy: 0.2231 - val_loss: 2.0459 - val_accuracy: 0.2279\n",
            "Epoch 70/75\n",
            "49/49 [==============================] - 4s 92ms/step - loss: 2.0783 - accuracy: 0.2248 - val_loss: 2.0393 - val_accuracy: 0.2284\n",
            "Epoch 71/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.0726 - accuracy: 0.2258 - val_loss: 2.0355 - val_accuracy: 0.2311\n",
            "Epoch 72/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.0702 - accuracy: 0.2280 - val_loss: 2.0316 - val_accuracy: 0.2305\n",
            "Epoch 73/75\n",
            "49/49 [==============================] - 4s 90ms/step - loss: 2.0671 - accuracy: 0.2290 - val_loss: 2.0278 - val_accuracy: 0.2300\n",
            "Epoch 74/75\n",
            "49/49 [==============================] - 4s 90ms/step - loss: 2.0636 - accuracy: 0.2303 - val_loss: 2.0239 - val_accuracy: 0.2315\n",
            "Epoch 75/75\n",
            "49/49 [==============================] - 4s 91ms/step - loss: 2.0618 - accuracy: 0.2302 - val_loss: 2.0220 - val_accuracy: 0.2311\n",
            "Network takes 366.261 seconds to train\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AGq92F3mtPOC",
        "colab_type": "code",
        "outputId": "bce8ddac-f335-4199-d9bd-e060188e0585",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 68
        }
      },
      "source": [
        "model.save_weights(\"resnet_bn_only_2.h5\")\n",
        "!gsutil -m cp -r resnet_bn_only_2.h5 gs://batch_norm_tpu/"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Copying file://resnet_bn_only_2.h5 [Content-Type=application/octet-stream]...\n",
            "/ [1/1 files][  2.4 MiB/  2.4 MiB] 100% Done                                    \n",
            "Operation completed over 1 objects/2.4 MiB.                                      \n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}